// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <melon/container/f14_map.h>
#include <pollux/vector/flat_vector.h>

namespace kumo::pollux {
    /// A thread-level cache of pre-allocated flat vectors of different types.
    /// Keeps up to 10 recyclable vectors of each type. A vector is
    /// recyclable if it is flat and recursively singly-referenced.
    /// Only singleton built-in types are supported. Decimal types, fixed-size array
    /// type, complex and custom types are not supported. Calling 'get' for an
    /// unsupported type already returns a newly allocated vector. Calling 'release'
    /// for an unsupported type is a no-op.
    class VectorPool {
    public:
        explicit VectorPool(memory::MemoryPool *pool) : pool_{pool} {
        }

        /// Gets a possibly recycled vector of 'type and 'size'. Allocates from
        /// 'pool_' if no pre-allocated vector or type is a complex type.
        VectorPtr get(const TypePtr &type, vector_size_t size);

        /// Moves vector into 'this' if it is flat, recursively singly referenced and
        /// there is space. The function returns true if 'vector' is not null and has
        /// been returned back to this pool, otherwise returns false.
        bool release(VectorPtr &vector);

        size_t release(std::vector<VectorPtr> &vectors);

        /// Clears all the cached vectors.
        void clear();

    private:
        /// Max number of elements for a vector to be recyclable. The larger
        /// the batch the less the win from recycling.
        static constexpr vector_size_t kMaxRecycleSize = 64 * 1024;
        static constexpr int32_t kNumPerType = 10;

        struct TypePool {
            int32_t size{0};
            std::array<VectorPtr, kNumPerType> vectors;

            bool maybePushBack(VectorPtr &vector);

            VectorPtr pop(
                const TypePtr &type,
                vector_size_t vectorSize,
                memory::MemoryPool &pool);

            /// Clears all the cached vectors.
            void clear();
        };

        memory::MemoryPool *const pool_;

        static constexpr int32_t kNumCachedVectorTypes =
                static_cast<int32_t>(TypeKind::HUGEINT) + 1;

        /// Caches of pre-allocated vectors indexed by typeKind.
        std::array<TypePool, kNumCachedVectorTypes> vectors_;
    };

    /// A simple vector ptr wrapper with an associated vector pool. It releases
    /// the allocated vector back to vector pool on destruction.
    class VectorRecycler {
    public:
        VectorRecycler(VectorPtr &vector, VectorPool *pool)
            : pool_(pool), vector_(vector) {
        }

        VectorRecycler(const VectorRecycler &) = delete;

        VectorRecycler &operator=(const VectorRecycler &) = delete;

        VectorRecycler(const VectorRecycler &&) = delete;

        VectorRecycler &operator=(const VectorRecycler &&) = delete;

        VectorRecycler(VectorRecycler &) = delete;

        VectorRecycler &operator=(VectorRecycler &) = delete;

        ~VectorRecycler() {
            if (pool_) {
                pool_->release(vector_);
            }
        }

    private:
        VectorPool *const pool_;
        VectorPtr &vector_;
    };
} // namespace kumo::pollux
